import os
import random
import shutil
from collections import defaultdict


def split_files(path, specified_string_groups):
    """
    Split files into train, val, and test sets based on specified string groups.

    Args:
        path (str): The directory path containing the files.
        specified_string_groups (list): List of lists of strings to specify groups.

    Returns:
        tuple: Lists of train files, val files, and test files.
    """
    # 获取指定路径下的所有文件名
    file_names = os.listdir(path)
    # 过滤掉文件夹，只保留文件
    file_names = [f for f in file_names if os.path.isfile(os.path.join(path, f))]
    
    # 按文件名开头分组
    file_groups = defaultdict(list)
    for file in file_names:
        for group_keys in specified_string_groups:
            if any(file.startswith(group_key) for group_key in group_keys):
                group_key = '_'.join(group_keys)  # Create a unique key for the group
                file_groups[group_key].append(file)
                break
    
    # 打乱文件组列表，以保证随机性
    group_keys = list(file_groups.keys())
    random.shuffle(group_keys)
    
    total_groups = len(group_keys)
    # 计算训练集、验证集和测试集的文件组数量
    train_val_count = int(total_groups * 0.85)
    test_count = total_groups - train_val_count
    
    # 分割文件组列表
    train_val_groups = group_keys[:train_val_count]
    test_groups = group_keys[train_val_count:]
    
    # 将文件组转换为文件列表
    train_files = [file for group in train_val_groups for file in file_groups[group]]
    random.shuffle(train_files)
    # val_files = []  # 可以根据需要进一步分割 train_files 为 train 和 val
    val_files = train_files[:int(len(train_files) * 0.15)]  # 分割 train_files 为 train 和 val
    train_files = train_files[int(len(train_files) * 0.15):]
    test_files = [file for group in test_groups for file in file_groups[group]]

    txtpath = path.replace('images', 'trainvalsettxt')
    if not os.path.exists(txtpath):  # 使用 os.path.exists() 检查路径是否存在
        os.makedirs(txtpath)  # 若路径不存在，则尝试创建多级目录
    
    # 写入训练集文件
    with open(os.path.join(txtpath, 'train.txt'), 'w') as train_file:
        for file in train_files:
            train_file.write(file + '\n')

    # 写入验证集文件
    with open(os.path.join(txtpath, 'val.txt'), 'w') as val_file:
        for file in val_files:
            val_file.write(file + '\n')

    # 写入测试集文件
    with open(os.path.join(txtpath, 'test.txt'), 'w') as test_file:
        for file in test_files:
            test_file.write(file + '\n')

    return train_files, val_files, test_files


def move_files(path, train_files, val_files, test_files):
    # 创建 train、val、test 文件夹
    train_dir = os.path.join(path, 'train')
    val_dir = os.path.join(path, 'val')
    test_dir = os.path.join(path, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 移动训练集文件
    for file in train_files:
        src = os.path.join(path, file)
        dst = os.path.join(train_dir, file)
        shutil.move(src, dst)

    # 移动验证集文件
    for file in val_files:
        src = os.path.join(path, file)
        dst = os.path.join(val_dir, file)
        shutil.move(src, dst)

    # 移动测试集文件
    for file in test_files:
        src = os.path.join(path, file)
        dst = os.path.join(test_dir, file)
        shutil.move(src, dst)

def gen_images_set_and_move_file(path, specified_string_groups):
    train_files, val_files, test_files = split_files(path, specified_string_groups)
    move_files(path, train_files, val_files, test_files)

def move_files_based_on_txt(path,txt1,txt2,txt3):
    # 创建 train、val、test 文件夹
    train_dir = os.path.join(path, 'train')
    val_dir = os.path.join(path, 'val')
    test_dir = os.path.join(path, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 读取 train.txt 并移动文件
    with open(txt1, 'r') as f:
        train_files = [os.path.splitext(line.strip())[0] for line in f.readlines()]
        for file in os.listdir(path):
            file_name = os.path.splitext(file)[0]
            if file_name in train_files:
                src = os.path.join(path, file)
                dst = os.path.join(train_dir, file)
                shutil.move(src, dst)

    # 读取 val.txt 并移动文件
    with open(txt2, 'r') as f:
        val_files = [os.path.splitext(line.strip())[0] for line in f.readlines()]
        for file in os.listdir(path):
            file_name = os.path.splitext(file)[0]
            if file_name in val_files:
                src = os.path.join(path, file)
                dst = os.path.join(val_dir, file)
                shutil.move(src, dst)

    # 读取 test.txt 并移动文件
    with open(txt3, 'r') as f:
        test_files = [os.path.splitext(line.strip())[0] for line in f.readlines()]
        for file in os.listdir(path):
            file_name = os.path.splitext(file)[0]
            if file_name in test_files:
                src = os.path.join(path, file)
                dst = os.path.join(test_dir, file)
                shutil.move(src, dst)
if __name__ == "__main__":
    # 1. 首先注释下面的代码，只是gen_images_set_and_move_file
    # image_path = "data_zoo/splitsynbboxpcb/images"
    image_path="data_zoo/splittruecolorwithdepthseg/images"
    specified_string_groups=[['light_img_0_0','light_img_0_1'],['light_img_0_2','light_img_0_3'],['light_img_0_4','light_img_0_5'],['light_img_0_6','light_img_0_7'],['light_img_10']]
    # gen_images_set_and_move_file(image_path,specified_string_groups)


    # 2. 然后选择下面的特定path，运行move_files_based_on_txt
    # label_path="data_zoo/splitsynbboxpcb/labels"
    # txt1="data_zoo/splitsynbboxpcb/trainvalsettxt/train.txt"
    # txt2="data_zoo/splitsynbboxpcb/trainvalsettxt/val.txt"
    # txt3="data_zoo/splitsynbboxpcb/trainvalsettxt/test.txt"


    # image_path = "data_zoo/pcbsubcp4clsbbg640/images"
    


    # label_path="data_zoo/pcbsubcp4clsbbg640/labels"
    # txt1="data_zoo/pcbsubcp4clsbbg640/trainvalsettxt/train.txt"
    # txt2="data_zoo/pcbsubcp4clsbbg640/trainvalsettxt/val.txt"
    # txt3="data_zoo/pcbsubcp4clsbbg640/trainvalsettxt/test.txt"


    # image_path = "data_zoo/splittruecolorseg/images"
    # label_path="data_zoo/splittruecolorseg/labels"
    # bbox_path="data_zoo/splittruecolorseg/bboxs"
    # mask_path="data_zoo/splittruecolorseg/mask"
    # txt1="data_zoo/splittruecolorseg/trainvalsettxt/train.txt"
    # txt2="data_zoo/splittruecolorseg/trainvalsettxt/val.txt"
    # txt3="data_zoo/splittruecolorseg/trainvalsettxt/test.txt"
    image_path = "data_zoo/splittruecolorwithdepthseg/images"
    label_path="data_zoo/splittruecolorwithdepthseg/labels"
    bbox_path="data_zoo/splittruecolorwithdepthseg/bboxs"
    mask_path="data_zoo/splittruecolorwithdepthseg/mask"
    txt1="data_zoo/splittruecolorwithdepthseg/trainvalsettxt/train.txt"
    txt2="data_zoo/splittruecolorwithdepthseg/trainvalsettxt/val.txt"
    txt3="data_zoo/splittruecolorwithdepthseg/trainvalsettxt/test.txt"


    txtpath=label_path.replace('labels','trainvalsettxt')
    if not os.path.exists(txtpath):  # 使用 os.path.exists() 检查路径是否存在
        os.makedirs(txtpath)  # 若路径不存在，则尝试创建多级目录
    # gen_images_set_and_move_file(image_path)
    move_files_based_on_txt(label_path,txt1,txt2,txt3)


    #TLDR
    """
    先执行gen_images_set_and_move_file(image_path)注释move_files_based_on_txt(label_path,txt1,txt2,txt3)
    最后注释gen_images_set_and_move_file执行move_files_based_on_txt
    """
    

    